/******************************************************************************/
/* @ All Rights Reserved Copyright (C) 2014 broada,Ltd.                       */
/*                                                                            */
/*  NAME      = postgresql_core.c                                             */
/*  NOTE      = POSTGRESQL协议解析                                                  */
/*  DATE      = 2015/03/18 by zhengxd                                         */
/******************************************************************************/
#include <stdio.h>
#include <string.h>
#include <netinet/in.h>

#include "frservice_comm.h"
#include "postgresql_core.h"
#include "postgresql_comm.h"
#include "postgresql_request.h"
#include "postgresql_response.h"
#include "config.h"

#define POSTGRESQL_HASH_COUNT          100           /* 最大hash个数 */
#define POSTGRESQL_MAX_BUFF_SIZE       1024 * 1024

extern void *g_p_apm_context;

/*******************************************************************************
*  FUNC     :  初始化私有postgresql session
*  ARGS     :  p_st_session(POSTGRESQL私有session)
*  RTN      :  FR_SUCC       (normal)
*              FR_FAIL       (error)
*  NOTE     :  p_st_session必须不为NULL
*******************************************************************************/
FR_STATIC int __postgresql_init_private(L7_session *p_st_session)
{
    Frhash_header *p_st_prepare_hash = NULL;                    /* prepare hash */
    L7_fragment_postgresql *p_st_fragment = NULL;               /* 分片包相关数据 */
    L7session_postgresql  *p_st_postgresql_conn = NULL;         /* 链接信息 */
    L7session_postgresql_private *p_st_private_session = NULL;  /* POSTGRESQL私有session */

    /* 申请整个私有session结构的内存 */
    p_st_private_session = (L7session_postgresql_private *)frmalloc(sizeof(L7session_postgresql_private));
    if (p_st_private_session == NULL)
    {
        goto error;
    }
    memset(p_st_private_session, 0x00, sizeof(L7session_postgresql_private));

    /* 申请链路信息结构内存 */
    p_st_postgresql_conn = (L7session_postgresql *)frmalloc(sizeof(L7session_postgresql));
    if (p_st_postgresql_conn == NULL)
    {
        goto error;
    }
    memset(p_st_postgresql_conn, 0x00, sizeof(L7session_postgresql));

    /* 申请分片包信息结构内存 */
    p_st_fragment = (L7_fragment_postgresql *)frmalloc(sizeof(L7_fragment_postgresql));
    if (p_st_fragment == NULL)
    {
        goto error;
    }
    memset(p_st_fragment, 0x00, sizeof(L7_fragment_postgresql));
    p_st_fragment->i_need_req_resp = FR_CLIENT_IP_PORT;

    /* 创建prepare hash */
    p_st_prepare_hash = frhash2_create2(POSTGRESQL_HASH_COUNT);
    if (p_st_prepare_hash == NULL)
    {
        goto error;
    }

    postgresql_quicksort(0, -1);
    /* 设置初始化结构 */
    p_st_private_session->p_st_sess = p_st_postgresql_conn;
    p_st_private_session->p_st_fragment = p_st_fragment;
    p_st_private_session->p_hash_prepare = p_st_prepare_hash;
    p_st_session->p_session_private = p_st_private_session;
    return FR_SUCC;

error:
    if (p_st_private_session)
    {
        frfree(p_st_private_session);
    }
    if (p_st_postgresql_conn)
    {
        frfree(p_st_postgresql_conn);
    }
    if (p_st_fragment)
    {
        frfree(p_st_fragment);
    }
    return FR_FAIL;
}

/*******************************************************************************
*  FUNC     :  初始化私有postgresql session
*  ARGS     :  p_st_private(私有session数据)
*  RTN      :  FR_SUCC       (normal)
*              FR_FAIL       (error)
*  NOTE     :  p_st_private必须不是NULL
*******************************************************************************/
FR_STATIC void __postgresql_destroy_curr_info(L7session_postgresql_private *p_st_private)
{
    if (p_st_private->p_st_req)
    {
        postgresql_destroy_request(&p_st_private->p_st_req);
    }
    if (p_st_private->p_st_resp)
    {
        postgresql_destroy_response(&p_st_private->p_st_resp);
    }
    postgresql_init_fragment(p_st_private->p_st_fragment, FR_CLIENT_IP_PORT);

    return;
}

/*******************************************************************************
*  FUNC     :  取得完整的packet包
*  ARGS     :  p_st_conn(connection数据)
*           :  p_st_packet(链路数据)
*           :  p_st_private(私有session数据)
*  RTN      :  APM_L7_ANALYZE_FAIL (fail)
*              APM_L7_ANALYZE_TODO (need more data)
*              APM_L7_ANALYZE_SUCC (succ)
*  NOTE     :
*******************************************************************************/
FR_STATIC int __postgresql_set_analysis_data(Conntrack *p_st_conn
                                          , Apm_packet *p_st_packet
                                          , L7session_postgresql_private *p_st_private)
{
    int i_ret = FR_SUCC;
    int i_req_resp = 0;
    unsigned int ui_len = p_st_packet->ui_payload_len;
    char *p_data = p_st_packet->p_payload;
    L7_fragment_postgresql *p_st_fragment = p_st_private->p_st_fragment;

    /* 无数据，跳过该包 */
    if (p_st_packet->ui_payload_len == 0)
    {
        return APM_L7_ANALYZE_TODO;
    }

    /* 取得方向 */
    i_req_resp = frcap_get_cs_direction(p_st_conn, p_st_packet->src_ip
                                       , p_st_packet->src_port);
    if (p_st_fragment->i_need_req_resp != i_req_resp)
    {
        frdbg("double request/response!\n");
        __postgresql_destroy_curr_info(p_st_private);
    }

    if (p_st_fragment->p_st_data == NULL)
    {
        p_st_fragment->p_st_data = bytebuff_creat();
        if (p_st_fragment->p_st_data == NULL)
        {
            return APM_L7_ANALYZE_FAIL;
        }
        i_ret = bytebuff_add(p_st_fragment->p_st_data, p_data, ui_len);
        if (i_ret != FR_SUCC)
        {
            return APM_L7_ANALYZE_FAIL;
        }

        p_st_fragment->i_req_resp = i_req_resp;
        p_st_fragment->st_start_timestamp = p_st_packet->st_timestamp;
        p_st_fragment->st_end_timestamp = p_st_packet->st_timestamp;
        return APM_L7_ANALYZE_SUCC;
    }
    else
    {
        /* 前后两次方向不一致, 消除上次信息 */
        if (p_st_fragment->i_req_resp != i_req_resp)
        {
            __postgresql_destroy_curr_info(p_st_private);
            i_ret = __postgresql_set_analysis_data(p_st_conn, p_st_packet, p_st_private);
            return i_ret;
        }
        else
        {
            i_ret = bytebuff_add(p_st_fragment->p_st_data, p_data, ui_len);
            if (i_ret != FR_SUCC)
            {
                return APM_L7_ANALYZE_FAIL;
            }
            p_st_fragment->st_end_timestamp = p_st_packet->st_timestamp;
            return APM_L7_ANALYZE_SUCC;
        }
    }
}

/*******************************************************************************
*  FUNC     :  判断当前数据链路是否是DRDA协议
*  ARGS     :  p_st_conn(入力链接信息)
*           :  p_st_packet(入力数据包)
*  RTN      :  APM_L7_IS_NO   (NO)
*              APM_L7_IS_YES  (YES)
*  NOTE     :
*******************************************************************************/
int postgresql_detect(Conntrack *p_st_conn, Apm_packet *p_st_packet)
{
    int i_ret = FR_SUCC;
    int i_req_resp = FR_DIRECT_UNKNOWN;
    uint8_t u8_type = 0;
    int i_packet_len = 0;
    Frbyte_Packet *p_st_data = NULL;
    unsigned int ui_len = 0;
    char *p_data = NULL;
    char *p_name = NULL;

    //入力参数检查
    if(p_st_packet == NULL || p_st_packet->ui_payload_len <= 4
       || p_st_packet->p_payload == NULL || p_st_conn == NULL)
    {
        return APM_L7_IS_NO;
    }
    ui_len = p_st_packet->ui_payload_len;        /* packet数据长度 */
    p_data = p_st_packet->p_payload;             /* packet数据内容 */

    p_st_data = bytebuff_creat();
    if (p_st_data == NULL)
    {
        frdbg("Create bytebuff failed!\n");
        return APM_L7_IS_NO;
    }
    i_ret = bytebuff_add(p_st_data, p_data, ui_len);
    if (i_ret != FR_SUCC)
    {
        frdbg("Add data to bytebuff failed!\n");
        bytebuff_free(&p_st_data);
        return APM_L7_IS_NO;
    }

    /* 取得第一个byte */
    i_ret = bytebuff_peek_char(p_st_data, &u8_type);
    if (i_ret != FR_SUCC)
    {
        frdbg("peek packet type failed!\n");
        goto error;
    }
    switch (u8_type)
    {
        case POSTGRESQL_TP_AUTHENTICATION:
        case POSTGRESQL_TP_BIND_COMPLETE:
        case POSTGRESQL_TP_CLOSE_COMPLETE:
        case POSTGRESQL_TP_EMPTY_QUERY_RESP:
        case POSTGRESQL_TP_FUNCTION_CALL_RESP:
        case POSTGRESQL_TP_NO_DATA://todo
        case POSTGRESQL_TP_NOTICE_RESP:
        case POSTGRESQL_TP_NOTIFICATION_RESP:
        case POSTGRESQL_TP_PARSE_COMPLETE:
        case POSTGRESQL_TP_PORTAL_SUSPENDED://todo
        case POSTGRESQL_TP_READY_QUERY://todo
        case POSTGRESQL_TP_ROW_DESCRIPTION:
        case POSTGRESQL_TP_BACKEND_KEY_DATA:
        case POSTGRESQL_TP_BIND:
        case POSTGRESQL_TP_FUNCTION_CALL:
        case POSTGRESQL_TP_PARSE:
        case POSTGRESQL_TP_PASSWORD_MSG:
        case POSTGRESQL_TP_QUERY:
        case POSTGRESQL_TP_CLOSE: //与POSTGRESQL_TP_COMMAND_COMPLETE相同
        case POSTGRESQL_TP_COPY_DATA:
        case POSTGRESQL_TP_COPY_DONE:
        case POSTGRESQL_TP_COPY_FAIL:
        case POSTGRESQL_TP_COPY_IN_RESP:
        case POSTGRESQL_TP_COPY_OUT_RESP://与POSTGRESQL_TP_FLUSH相同
        case POSTGRESQL_TP_DATA_ROW://与POSTGRESQL_TP_DESCRIBE相同
        case POSTGRESQL_TP_ERROR_RESP://与POSTGRESQL_TP_EXECUTE相同
        case POSTGRESQL_TP_PARAMETER_DESCRIPTION:
        case POSTGRESQL_TP_PARAMETER_STATUS://与POSTGRESQL_TP_SYNC相同
        case POSTGRESQL_TP_TERMINATE:
            break;

        default:
            /* 判断是否是Startup */
            i_ret = postgresql_get_int2(p_st_data, &i_packet_len);
            if (i_ret != FR_SUCC || i_packet_len < 8 || ui_len != (unsigned int)i_packet_len)
            {
                frdbg("Get Startup packet length failed!\n");
                goto error;
            }
            i_packet_len -= 8;
            /* skip protocol version number */
            bytebuff_skip(p_st_data, 4);

            while (i_packet_len)
            {
                i_ret = postgresql_get_string(p_st_data, &i_packet_len, &p_name);
                if (i_ret != FR_SUCC)
                {
                    frdbg("Get name failed!\n");
                    goto error;
                }
                /* 结束 */
                if (p_name == NULL)
                {
                    if (i_packet_len != 0)
                    {
                        frdbg("format incorrect!\n");
                        goto error;
                    }
                    /* 设置方向 */
                    bytebuff_free(&p_st_data);
                    if(frct_get_direct_known(p_st_conn) == FR_DIRECT_UNKNOWN)
                    {
                        i_ret = frcap_set_cs_srv(p_st_conn, p_st_packet->dst_ip
                                                 , p_st_packet->dst_port);
                        if (i_ret != FR_SUCC)
                        {
                            frdbg("Set server ip and port failed!\n");
                            return APM_L7_IS_NO;
                        }
                    }
                    frservice_add(g_p_apm_context, MODULE_L7PROTO
                                 , p_st_packet->dst_ip, p_st_packet->dst_port, NULL, 0);
                    return APM_L7_IS_YES;
                }
                frfree(p_name);
                p_name = NULL;
                /* skip value */
                i_ret = postgresql_skip_string(p_st_data, &i_packet_len);
                if (i_ret != FR_SUCC)
                {
                    frdbg("format incorrect!\n");
                    goto error;
                }
            }
    }

    while (bytebuff_remaining(p_st_data))
    {
        i_ret = bytebuff_get_char(p_st_data, &u8_type);
        if (i_ret != FR_SUCC)
        {
            frdbg("get packet type failed!\n");
            goto error;
        }
        switch (u8_type)
        {
            case POSTGRESQL_TP_BIND:
            case POSTGRESQL_TP_FUNCTION_CALL:
            case POSTGRESQL_TP_PARSE:
            case POSTGRESQL_TP_PASSWORD_MSG:
            case POSTGRESQL_TP_QUERY:
            //case POSTGRESQL_TP_CLOSE:
            case POSTGRESQL_TP_COPY_FAIL:
            //case POSTGRESQL_TP_FLUSH:
            //case POSTGRESQL_TP_DESCRIBE:
            //case POSTGRESQL_TP_EXECUTE:
            //case POSTGRESQL_TP_SYNC:
            case POSTGRESQL_TP_TERMINATE:
                i_req_resp = FR_CLIENT_IP_PORT;
                break;

            case POSTGRESQL_TP_AUTHENTICATION:
            case POSTGRESQL_TP_BACKEND_KEY_DATA:
            case POSTGRESQL_TP_BIND_COMPLETE:
            case POSTGRESQL_TP_CLOSE_COMPLETE:
            //case POSTGRESQL_TP_COMMAND_COMPLETE:
            case POSTGRESQL_TP_COPY_IN_RESP:
            case POSTGRESQL_TP_COPY_OUT_RESP:
            case POSTGRESQL_TP_DATA_ROW:
            case POSTGRESQL_TP_EMPTY_QUERY_RESP:
            //case POSTGRESQL_TP_ERROR_RESP:
            case POSTGRESQL_TP_FUNCTION_CALL_RESP:
            case POSTGRESQL_TP_NO_DATA:
            case POSTGRESQL_TP_NOTICE_RESP:
            case POSTGRESQL_TP_NOTIFICATION_RESP:
            case POSTGRESQL_TP_PARAMETER_DESCRIPTION:
            //case POSTGRESQL_TP_PARAMETER_STATUS:
            case POSTGRESQL_TP_PARSE_COMPLETE:
            case POSTGRESQL_TP_PORTAL_SUSPENDED:
            case POSTGRESQL_TP_READY_QUERY:
            case POSTGRESQL_TP_ROW_DESCRIPTION:
                i_req_resp = FR_SRV_IP_PORT;
                break;

            case POSTGRESQL_TP_COPY_DATA:
            case POSTGRESQL_TP_COPY_DONE:
                break;

            default:
                frdbg("unknow packet type!\n");
                goto error;
         }

        i_ret = postgresql_get_int2(p_st_data, &i_packet_len);
        if (i_ret != FR_SUCC || i_packet_len < 4
            || i_packet_len - 4 > bytebuff_remaining(p_st_data))
        {
            frdbg("packet length not match!\n");
            goto error;
        }
        bytebuff_skip(p_st_data, i_packet_len - 4);
    }

    /* 设置方向 */
    if(frct_get_direct_known(p_st_conn) == FR_DIRECT_UNKNOWN)
    {
        if (i_req_resp == FR_SRV_IP_PORT)
        {
            i_ret = frcap_set_cs_srv(p_st_conn, p_st_packet->src_ip
                                     , p_st_packet->src_port);
        }
        else if (i_req_resp == FR_CLIENT_IP_PORT)
        {
            i_ret = frcap_set_cs_srv(p_st_conn, p_st_packet->dst_ip
                                            , p_st_packet->dst_port);
        }
        else
        {
            i_ret = FR_FAIL;
        }
    }

    if (i_ret == FR_FAIL)
    {
        goto error;
    }
    bytebuff_free(&p_st_data);
    frservice_add(g_p_apm_context, MODULE_L7PROTO, p_st_packet->dst_ip, p_st_packet->dst_port, NULL, 0);
    return APM_L7_IS_YES;

error:
    bytebuff_free(&p_st_data);
    /* 有可能方向未知 */
    i_ret = frservice_get(g_p_apm_context, MODULE_L7PROTO, p_st_packet->dst_ip, p_st_packet->dst_port, NULL, 0, NULL);
    if (i_ret == FR_SUCC)
    {
        frcap_set_cs_srv(p_st_conn, p_st_packet->dst_ip, p_st_packet->dst_port);
        return APM_L7_IS_YES;
    }
    i_ret = frservice_get(g_p_apm_context, MODULE_L7PROTO, p_st_packet->src_ip, p_st_packet->src_port, NULL, 0, NULL);
    if (i_ret == FR_SUCC)
    {
        frcap_set_cs_srv(p_st_conn, p_st_packet->src_ip, p_st_packet->src_port);
        return APM_L7_IS_YES;
    }
    return APM_L7_IS_NO;
}

/*******************************************************************************
*  FUNC     :  POSTGRESQL的packet解析
*  ARGS     :  p_st_conn(入力链接信息)
            :  p_st_packet(入力packet信息)
            :  p_st_hash_report(保存Frreport_data)
*  RTN      :  参照module.h
*  NOTE     :
*******************************************************************************/
int postgresql_analysis(Conntrack *p_st_conn
                        , Apm_packet *p_st_packet
                        , Frhash_header *p_st_hash_report)
{
    int i_ret = FR_SUCC;
    L7_session *p_st_session = NULL;
    L7session_postgresql_private *p_st_private = NULL;

    if(p_st_conn == NULL || p_st_packet == NULL || p_st_hash_report == NULL)
    {
        return APM_L7_ANALYZE_FAIL;
    }

    /* 到Conntrack中取得session信息 */
    p_st_session = frct_get_session(p_st_conn);
    if (p_st_session == NULL)
    {
        return APM_L7_ANALYZE_FAIL;
    }

    /*
    * 如果入力p_st_packet是FIN包, 说明本次POSTGRESQL的connection结束,
    * 释放本次session信息. 如果入力p_st_packet是SYN包,
    * 说明新的POSTGRESQL的connection开始了, 释放上次可能残留的session信息
    */
    if (p_st_packet->em_analyz_cmd == FRANALYZ_CMD_FIN
        || p_st_packet->em_analyz_cmd == FRANALYZ_CMD_SYN)
    {
        /* 销毁所有相关本次POSTGRESQL的connection的数据 */
        postgresql_session_destroy((void *)(p_st_session->p_session_private));
        p_st_session->p_session_private = NULL;
        return APM_L7_ANALYZE_SUCC;
    }

    /* 本次链路的第一个有效packet包,创建链路 */
    if (p_st_session->p_session_private == NULL)
    {
        /* 初始化POSTGRESQL的私有数据 */
        i_ret = __postgresql_init_private(p_st_session);
        if (i_ret != FR_SUCC)
        {
            return APM_L7_ANALYZE_FAIL;
        }
        /* 设置session ID */
        fruuid(p_st_session->arr_session_id);
    }

    p_st_private = (L7session_postgresql_private *)p_st_session->p_session_private;
    /* 设置解析数据 */
    i_ret = __postgresql_set_analysis_data(p_st_conn, p_st_packet, p_st_private);
    if (i_ret == APM_L7_ANALYZE_TODO)
    {
        return APM_L7_ANALYZE_TODO;
    }
    else if (i_ret == APM_L7_ANALYZE_FAIL)
    {
        frdbg("call __postgresql_set_analysis_data failed\n");
        __postgresql_destroy_curr_info(p_st_private);
        return APM_L7_ANALYZE_FAIL;
    }

    /* 解析packet */
    switch (p_st_private->p_st_fragment->i_req_resp)
    {
        case FR_CLIENT_IP_PORT: /* 解析request */
            i_ret = postgresql_parse_request(p_st_conn
                                            , p_st_hash_report
                                            , p_st_private);
            if (i_ret == APM_L7_ANALYZE_FAIL)
            {
                __postgresql_destroy_curr_info(p_st_private);
                return APM_L7_ANALYZE_FAIL;
            }
            break;

        case FR_SRV_IP_PORT: /* 解析response */
            i_ret = postgresql_parse_response(p_st_conn, p_st_hash_report
                                             , p_st_private);
            if (i_ret == APM_L7_ANALYZE_TODO)
            {
                return APM_L7_ANALYZE_TODO;
            }
            __postgresql_destroy_curr_info(p_st_private);
            break;

         default:
            frdbg("unkonow\n");
            i_ret = APM_L7_ANALYZE_FAIL;
    }

    /* 返回 */
    return i_ret;
}

/*******************************************************************************
*  FUNC     :  私有信息销毁
*  ARGS     :  p_private_session(POSTGRESQL的私有session数据)
*  RTN      :  void
*  NOTE     :
*******************************************************************************/
int postgresql_session_destroy(void *p_private_session)
{
    L7session_postgresql_private *p_st_private = (L7session_postgresql_private *)p_private_session;

    /* 判断session是否为空 */
    if (p_st_private == NULL)
    {
        return FR_SUCC;
    }

    /* 销毁session数据 */
    if (p_st_private->p_st_sess)
    {
        postgresql_destroy_session(&(p_st_private->p_st_sess));
    }

    /* 销毁request数据 */
    if (p_st_private->p_st_req)
    {
        postgresql_destroy_request(&(p_st_private->p_st_req));
    }

    /* 销毁response数据 */
    if (p_st_private->p_st_resp)
    {
        postgresql_destroy_response(&p_st_private->p_st_resp);
    }

    /* 销毁分片数据 */
    if (p_st_private->p_st_fragment)
    {
        postgresql_destroy_fragment(&p_st_private->p_st_fragment);
    }

    /* 销毁prepare数据 */
    frhash2_free2(p_st_private->p_hash_prepare, frdb_statement_node_free);

    frfree(p_st_private);

    return FR_SUCC;
}
